# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment4/question1/main_plot.py
# Modified from https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment3/question2/main_plot.py
import matplotlib.pyplot
import mpl_toolkits.mplot3d
import torch
import torch.utils.data

import datasets
import models


def main():
    model = models.LorentzDispersionModel()
    model.load_state_dict(torch.load("./outputs/trained_model.pt"))

    dataset = datasets.EpsDataset()
    data_loader = torch.utils.data.DataLoader(dataset, batch_size=len(dataset))
    data_x, data_y = next(iter(data_loader))

    model_x = torch.linspace(float(torch.min(data_x)), float(torch.max(data_x)), 1000)
    model_y = torch.detach(model(model_x))

    matplotlib.pyplot.subplot(221)
    matplotlib.pyplot.title("Real Part")
    matplotlib.pyplot.scatter(data_x, data_y[:, 0], c="lavender")
    matplotlib.pyplot.plot(model_x, model_y[:, 0], c="orange")
    matplotlib.pyplot.legend(["data point", "fitted curve"])
    matplotlib.pyplot.xlabel(r"$\omega$")
    # noinspection SpellCheckingInspection
    matplotlib.pyplot.ylabel(r"$\epsilon$")

    matplotlib.pyplot.subplot(222)
    matplotlib.pyplot.title("Imaginary Part")
    matplotlib.pyplot.scatter(data_x, data_y[:, 1], c="lavender")
    matplotlib.pyplot.plot(model_x, model_y[:, 1], c="orange")
    matplotlib.pyplot.xlabel(r"$\omega$")
    # noinspection SpellCheckingInspection
    matplotlib.pyplot.ylabel(r"$\epsilon$")

    matplotlib.pyplot.subplot(223)
    matplotlib.pyplot.title("Absolute")
    matplotlib.pyplot.scatter(
        data_x,
        torch.sqrt(data_y[:, 0] ** 2 + data_y[:, 1] ** 2),
        c="lavender")
    matplotlib.pyplot.plot(
        model_x,
        torch.sqrt(model_y[:, 0] ** 2 + model_y[:, 1] ** 2),
        c="orange")
    matplotlib.pyplot.xlabel(r"$\omega$")
    # noinspection SpellCheckingInspection
    matplotlib.pyplot.ylabel(r"$\epsilon$")

    axes: mpl_toolkits.mplot3d.Axes3D = matplotlib.pyplot.subplot(224, projection="3d")
    matplotlib.pyplot.title("3D")
    axes.scatter(data_x, data_y[:, 0], data_y[:, 1], c="lavender", alpha=0.8)
    axes.plot(model_x, model_y[:, 0], model_y[:, 1], c="orange")
    matplotlib.pyplot.xlabel(r"$\omega$")
    # noinspection SpellCheckingInspection
    matplotlib.pyplot.ylabel(r"$\epsilon_{\mathrm{Real}}$")
    # noinspection SpellCheckingInspection
    axes.set_zlabel(r"$\epsilon_{\mathrm{Imag}}$")

    matplotlib.pyplot.subplots_adjust(
        left=0.071, bottom=0.079, right=0.96, top=0.93,
        wspace=0.2, hspace=0.31)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
